import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from matplotlib.font_manager import FontProperties

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

# 物理参数
u0 = 100.0  # 初始温度 (℃)
l = 1.0     # 杆的长度 (m)
a = 0.1     # 热扩散系数 (m²/s)

# 空间和时间设置
x = np.linspace(0, l, 200)  # 空间网格
t_max = 20.0                # 最大时间 (s)
t = np.linspace(0, t_max, 200)  # 时间网格

# 计算温度分布函数
def temperature(x, t, N_terms=50):
    """
    计算杆上温度分布
    
    参数:
    x: 空间位置数组
    t: 时间
    N_terms: 级数求和的项数
    
    返回:
    u: 温度分布数组
    """
    u = np.zeros_like(x)
    for p in range(N_terms):
        coef = (-1)**p / (p + 0.5)**2
        exponent = -((p + 0.5) * np.pi * a / l)**2 * t
        sine = np.sin((p + 0.5) * np.pi * x / l)
        u += coef * np.exp(exponent) * sine
    
    u *= (2 * u0) / (np.pi**2)
    return u

# 创建图形和坐标轴
fig, ax = plt.subplots(figsize=(10, 6))
ax.set_xlim(0, l)
ax.set_ylim(-5, u0 + 10)
ax.grid(True, linestyle='--', alpha=0.7)
ax.set_xlabel('位置 x (m)', fontsize=12)
ax.set_ylabel('温度 u(x,t) (℃)', fontsize=12)
ax.set_title('细杆导热问题 - 温度分布随时间变化', fontsize=14)

# 初始线
line, = ax.plot([], [], 'r-', linewidth=2, label=f'当前温度分布')

# 添加固定点和绝热点标记
ax.plot(0, 0, 'bo', markersize=8, label='固定温度点 (0℃)')
ax.plot(l, 0, 'go', markersize=8, label='绝热点 ($\partial u/\partial x=0$)')

# 添加初始温度分布
initial_temp = u0 * x / l
ax.plot(x, initial_temp, 'k--', alpha=0.5, label='初始温度分布')

# 添加稳态线 (t→∞ 时 u→0)
ax.axhline(0, color='purple', linestyle='-', alpha=0.3, label='稳态 (0℃)')

# 添加图例
ax.legend(loc='upper right', fontsize=10)

# 添加时间文本
time_text = ax.text(0.02, 0.95, '', transform=ax.transAxes, fontsize=12)

# 动画初始化函数
def init():
    line.set_data([], [])
    time_text.set_text('')
    return line, time_text

# 动画更新函数
def update(frame):
    # 计算当前时间对应的温度分布
    current_t = t[frame]
    u = temperature(x, current_t)
    
    # 更新曲线
    line.set_data(x, u)
    
    # 更新时间文本
    time_text.set_text(f'时间: t = {current_t:.2f} s')
    
    return line, time_text

# 创建动画
ani = FuncAnimation(fig, update, frames=len(t),
                    init_func=init, blit=True, interval=50)

# 添加物理参数说明
param_text = ax.text(0.65, 0.85, 
                     f'物理参数:\n$u_0$ = {u0} ℃\n$l$ = {l} m\n$a^2$ = {a**2:.4f} m²/s',
                     transform=ax.transAxes, 
                     bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.5),
                     fontsize=10)

# 保存动画为GIF
ani.save('热传导8-1例题2.gif', writer='pillow', fps=20, dpi=100)

plt.tight_layout()
plt.show()